package ai.nlp;

import ai.djl.huggingface.tokenizers.Encoding;
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
import ai.djl.inference.Predictor;
import ai.djl.modality.Classifications;
import ai.djl.ndarray.NDList;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import org.apache.shardingsphere.extension.common.util.collection.ListUtil;

import java.util.List;

/**
 * 运行失败，异常:ai.djl.translate.TranslateException: java.lang.IllegalStateException: DataType mismatch, Required double Actual float32
 */
public class BertExample {

    public static void main(String[] args) {
        String text = "This movie was absolutely fantastic!";

        try (ZooModel<String, Classifications> model = loadModel();
             Predictor<String, Classifications> predictor = model.newPredictor()) {

            Classifications result = predictor.predict(text);

            // 遍历分类结果（修正 forEach 用法）
            System.out.println("分类结果:");
            for (Classifications.Classification classification : result.items()) {
                System.out.printf("%s: %.4f%n",
                        classification.getClassName(),
                        classification.getProbability());
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static ZooModel<String, Classifications> loadModel() throws Exception {
        // 使用自定义 Translator（替代旧的 TextClassificationTranslator）
        Translator<String, Classifications> translator = new BertTranslator();

        Criteria<String, Classifications> criteria = Criteria.builder()
                .setTypes(String.class, Classifications.class)
                .optModelUrls("djl://ai.djl.huggingface.pytorch/bert-base-uncased")
                .optTranslator(translator)
                .optEngine("PyTorch")
                .build();

        return criteria.loadModel();
    }

    // 自定义 Translator 实现
    static class BertTranslator implements Translator<String, Classifications> {

        private HuggingFaceTokenizer tokenizer;

        @Override
        public void prepare(TranslatorContext ctx) {
            // 初始化分词器
            tokenizer = HuggingFaceTokenizer.newInstance("bert-base-uncased");
        }

        @Override
        public NDList processInput(TranslatorContext ctx, String input) {
            // 分词并转换为张量
            Encoding encoding = tokenizer.encode(input);
            return new NDList(
                    ctx.getNDManager().create(encoding.getIds()),     // input_ids
                    ctx.getNDManager().create(encoding.getAttentionMask())     // attention_mask
            );
        }

        @Override
        public Classifications processOutput(TranslatorContext ctx, NDList ndList) {
            // 解析模型输出（假设输出为 [batch_size, num_labels]）
            float[] logits = ndList.get(0).toFloatArray();
            List<String> es = ListUtil.newArrayList("LABEL_0", "LABEL_1");
            List<Double> list = ListUtil.newArrayList((double)logits[0], (double)logits[1]);
            return new Classifications(
                    es,  // 标签需根据实际模型调整
                    list
            );
        }
    }
}